import asyncio
import re
from typing import List, Optional

from pydantic import BaseModel

from metagpt.const import RESEARCH_PATH
from metagpt.logs import logger
from metagpt.roles.role import Role, RoleReactMode
from metagpt.schema import Message

from data_dyne_wrj.agent.agent_test1 import KeywordCollecter
from data_dyne_wrj.rag.kb_api import SimpleRelevantPaper

class Report(BaseModel):
    topic: str
    links: dict[str, list[str]] = None
    summaries: list[tuple[str, str]] = None
    content: str = ""

# class RelvantKeyword(BaseModel):
#     id: int
#     keyword: str
#     child_keywords: Optional[List["RelvantKeyword"]] = None
#     relvent_paper: Optional[List[SimpleRelevantPaper]] = None

#     def list_all_keywords(self):
#         keywords = [self.keyword]
#         if self.child_keywords:
#             for child in self.child_keywords:
#                 keywords.extend(child.list_all_keywords())
#         return keywords
    
class Survey(BaseModel):
    keywords: list[str]
    topic: str
    papers: List[SimpleRelevantPaper] = None
    content: str = ""


class SurveyResearcher(Role):
    name: str = "David"
    profile: str = "Survey Researcher"
    goal: str = "Gather information and conduct research for survey"
    constraints: str = "Ensure accuracy and relevance of information"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.set_actions([KeywordCollecter])
        self._set_react_mode(RoleReactMode.REACT.value, 4)
        if self.language not in ("en-us", "zh-cn"):
            logger.warning(f"The language `{self.language}` has not been tested, it may not work.")

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo
        msg = self.rc.memory.get(k=1)[0]
        if isinstance(msg.instruct_content, Report):
            instruct_content = msg.instruct_content
            topic = instruct_content.topic
        else:
            topic = msg.content

        if isinstance(todo, KeywordCollecter):
            keywords = await todo.run(topic)
            ret = Message(
                content="", instruct_content=None, role=self.profile, cause_by=todo
            )
        return ret